import torch
import torch.onnx
from torch import nn
from train import *


# 加载模型
model = get_model(10)
model.load_state_dict(torch.load('python/weight_14/best_model.pth'))

# 设置为评估模式
model.eval()

# 创建一个虚拟输入，其形状必须与模型输入形状相匹配
# 假设输入的形状是(batch_size, channels, height, width)
# 你需要根据你的模型来调整这个形状
x = torch.randn(8, 3, 16384, requires_grad=True)

# 导出模型
torch.onnx.export(model,               # 模型 being run
                  x,                   # 模型输入 (or a tuple for multiple inputs)
                  "best_model.onnx",   # 导出的文件名 (也可以是文件对象)
                  export_params=True,  # 是否存储模型参数
                  opset_version=10,    # ONNX版本
                  do_constant_folding=False,  # 是否执行常量折叠优化
                  input_names = ['input'],   # 输入名
                  output_names = ['output'], # 输出名
                  dynamic_axes={'input' : {0 : 'batch_size'},  # 指定动态尺寸
                                'output' : {0 : 'batch_size'}})

print("模型已成功转换为ONNX格式")
